import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from collections import OrderedDict


class MiniAlexNet(nn.Sequential):
    def __init__(self, outputs):
        super(MiniAlexNet, self).__init__()
        first = OrderedDict([
            ('conv0', nn.Conv2d(1, 32, kernel_size=12, stride=1, padding=0)),
            ('relu0', nn.ReLU(inplace=True)),
            ('maxPool0', nn.MaxPool2d(kernel_size=3, stride=2)),
            ('conv1', nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2)),
            ('relu1', nn.ReLU(inplace=True)),
            ('maxPool1', nn.MaxPool2d(kernel_size=3, stride=2)),
            # --- central ---
            ('conv2', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
            ('relu2', nn.ReLU(inplace=True)),
            ('conv3', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
            ('relu3', nn.ReLU(inplace=True)),
            ('conv4', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
            ('relu4', nn.ReLU(inplace=True)),
            ('maxPool2', nn.MaxPool2d(kernel_size=3, stride=2))
        ])
        
#         dense_central = OrderedDict([
#             ('conv2', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
#             ('relu2', nn.ReLU(inplace=True)),
#             ('conv3', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
#             ('relu3', nn.ReLU(inplace=True)),
#             ('conv4', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)),
#             ('relu4', nn.ReLU(inplace=True)),
#             ('maxPool2', nn.MaxPool2d(kernel_size=3, stride=2))
#         ])
        
        self.features = nn.Sequential(first)
#         self.features = first
         
        feature_size = 1600
        classifier =  [
            nn.Linear(feature_size, 1024),
            nn.Linear(1024, outputs)
        ]
        self.classifier = nn.Sequential(*classifier)
        
    def forward(self, x):
#         for s, layer in self.features.items():
#             x = layer(x)
#             print('%s -> %s' % (s, x.shape))
        x = self.features(x)
        out = F.relu(x, inplace=True)
        out = torch.flatten(out, 1)
        out = self.classifier(out)
        return out